""" Plot band structure from a VASP-EIGENCAR file.
The user-interface function is get_band().
This function requires three input files: POSCAR, KPOINTS, and EIGENVAL.
Below are the reasons:
- to calculate distances between k-points, it requires POSCAR file to 
  get the three basis vectors.
- to retrieve conveniently starting and end points of each segment, it
  requires KPOINTS file which contains user input.
- EIGENVAL file contains the data and is required.
GHB, 20180208.
"""
import sys
import numpy as np
from crystal import basis

def usage():
    """ Display a help message."""
    print("    This program reads 3 files: EIGENVAL, POSCAR, and KPOINTS.")


def read_basis(f="POSCAR"):
    """ Reads basis vectors from POSCAR."""
    f_POSCAR = open(f, "r")
    all_lines = f_POSCAR.readlines()
    f_POSCAR.close()

    a = float(all_lines[1])
    b = []
    for i in (2,3,4):
        b.extend([float(x)*a for x in all_lines[i].split()])
    return basis.Basis(b)


def read_band(E_fermi=0.0, f="EIGENVAL"):
    """ Reads eigen values from EIGENVAL file.
    
    Arguments:
    E_fermi: (float) all energies will be subtracted by E_fermi. It is often 
        required to align (i.e., offset down or up) the energies at a reference
        point. Default: 0.0.
    f: (string) file name of eigenvalues. Default: EIGENVAL.
    
    Returns:
    if spin polarized (number of spin components is 2): 
        (kc, band_e_up, band_e_down, occupation_up, occupation_down)
    if spin degenerated (number of spin components is 1):
        (kc, band_e, occupation)
    kc: distance of k-points relative to the first k-point.
    band_e_up, band_e_down, band_e: eigenvalues.
    """
    f_EIGENVAL = open(f, "r")
    all_lines = f_EIGENVAL.readlines()
    f_EIGENVAL.close()

    # The first line is: 
    #    number-of-atoms number-of-atoms-again 1 number-of-spin
    n_spin = int(all_lines[0].split()[3])

    # Line 5 contains two useful numbers: n_kpoints and n_bands.
    n_electron, n_kpoints, n_bands = [int(x) for x in all_lines[5].split()]
    kc = []

    if n_spin==1:
        band_e = []
        occupation = []
        for i in range(n_bands):  
            band_e.append([])
            occupation.append([])
        li = 7   # li means line_index
        for point in range(n_kpoints):
            kx, ky, kz,  weight= [float(x) for x in all_lines[li].split()]
            kc.append((kx, ky, kz))
            for i in range(n_bands):
                li += 1
                curr_energy, curr_occ = [float(x) for x in 
                    all_lines[li].split()[1:]]
                band_e[i].append(curr_energy-E_fermi)
                occupation[i].append(curr_occ)
            li += 2
        return (kc, band_e, occupation)
    elif n_spin==2:
        band_e_up = []
        band_e_down = []
        occupation_up = []
        occupation_down = []
        for i in range(n_bands):  
            band_e_up.append([])
            band_e_down.append([])
            occupation_up.append([])
            occupation_down.append([])
        li = 7   # li means line_index
        for point in range(n_kpoints):
            kx, ky, kz,  weight= [float(x) for x in all_lines[li].split()]
            kc.append((kx, ky, kz))
            for i in range(n_bands):
                li += 1
                e_up, e_down, o_up, o_down = [float(x) for x in 
                    all_lines[li].split()[1:]] 
                band_e_up[i].append(e_up-E_fermi)
                band_e_down[i].append(e_down-E_fermi)
                occupation_up[i].append(o_up)
                occupation_down[i].append(o_down)
            li += 2
        return (kc, band_e_up, band_e_down, occupation_up, occupation_down)
    else:
        print("Unexpected number of spin-components:", n_spin)
        print("This number has to be 1 or 2. EIGENVAL may be corrupted. Bye.")
        sys.exit(1)


def read_kc_from_KPOINTS(f="KPOINTS"):
    """ Reads reciprocal coordinates from line-mode KPOINTS."""
    f_KPOINTS = open(f, "r")
    all_lines = f_KPOINTS.readlines()
    f_KPOINTS.close()
    
    if all_lines[2][0] not in "lL":
        print("This routine is for reading k-path.")
        print("The KPOINTS file must be of the line mode (see VASP manual).")
        print("Misusing this routine is meaningless, and I will stop.")
        sys.exit(1)

    n_per_segment = int(all_lines[1])   # number of points per segment
    
    if all_lines[3][0] not in 'rR':
        print("This version only supports reciprocal coordinate system.")
        print("Cartesian coordinate system is not supported.")
        sys.exit(1)

    # here is the data structure for the line segments:
    # [ks0a, ks0b, ks1a, ks1b, ...]
    ks = []
    for current_line in all_lines[4:] :

        # check if current_line is empty
        temp_array = [float(x) for x in current_line.split()[0:3]]
        if len(temp_array) == 0:
            continue    # This is an empty line; skip it.
        ks.append(temp_array)

    if len(ks)%2 == 1:
        print("The k-path should contain even-number nodes, but it isn't so.")
        print("Please check whether KPOINTS is corrupted.")
        sys.exit(1)

    return [n_per_segment, ks]


def calc_kc_segments(n, k_coordinate, reciprocal_basis):
    """ Calculates distance of k-point segments.

    The main difficulty is that two neighboring segments may be disjoined,
    e.g. a two-segment line specified by
    [[0.5, 0.0, 0.0], [0.5, 0.5, 0.0], [0.0, 0.5, 0.5], [0, 0, 0]],
    for which the distance between [0.5, 0.5, 0.0] and [0.0, 0.5, 0.5] should
    be set to zero, even though they are different coordinates.
    """
    k0 = reciprocal_basis.to_cartesian(k_coordinate[0])
    dist_l = [0.0]   # dist_l is coordinates used for labels
    for i, p in enumerate(k_coordinate[1:]) :
        k1 = reciprocal_basis.to_cartesian(p)
        if i%2 ==0:   # the end point of a segment
            d = np.sqrt((k0[0]-k1[0])**2 + (k0[1]-k1[1])**2 + (k0[2]-k1[2])**2)
            dist_l.append(d + dist_l[int(i/2)])
      
        k0 = k1

    dist_c = []    # dist_c is for coordinates used for plot
    for i in range(1, len(dist_l)):
        dist_c.extend(np.linspace(dist_l[i-1], dist_l[i], n))
        
    return [np.array(dist_c), np.array(dist_l)]


def get_band(E_fermi=0.0, 
             f_POSCAR="POSCAR", 
             f_EIGENVAL="EIGENVAL",
             f_KPOINTS="KPOINTS"):
        
    """ Gets data for plotting band structures.

    This is the central function to be called by a user.
    Arguments: 
    E_fermi: an energy that will be subtracted from those input energies;
    f_POSCAR: file name of VASP's input structure, default POSCAR;
    f_EIGENVAL: file name of VASP's eigenvalues, default EIGENVAL;
    f_KPOINTS: file name of VASP's input k-points, default KPOINTS.

    Returns:
    [dist_c, band_e, dist_l, occupation] in which
        dist_c: culmulative distance between the k-points; This is x-axis.
        band_e: band energies. Len(band_e) is the number of bands.
        dist_l: culmulative distance for x-ticks.
        occupation: occupation of bands.
    """
    basis_reciprocal = read_basis(f_POSCAR).reciprocal()
    n_per_segments, kc2 = read_kc_from_KPOINTS(f_KPOINTS)
    dist_c, dist_l = calc_kc_segments(n_per_segments, kc2, basis_reciprocal)
    eigenval_data = read_band(E_fermi, f_EIGENVAL)

    if(len(eigenval_data)==3):
        kc, band_e, occupation = eigenval_data
        return (dist_c, band_e, dist_l, occupation)
    elif(len(eigenval_data)==5):
        kc, e_up, e_down, o_up, o_down = eigenval_data
        return (dist_c, [e_up, e_down], dist_l, [o_up, o_down])
    
if __name__ =="__main__":
    pass
